# -*- coding: utf-8 -*-
"""\
This file is part of the ptycho package.

Utility functions.
"""

__all__ = ['verbose', 'FType', 'CType', 'set_datatype', 'FFT', 'Fourier_projection']

import numpy as np
from pyE17 import io
io.h5options['UNSUPPORTED'] = 'ignore' # otherwise saving e.g. a function will fail. 
from pyE17 import verbose
import scipy
import scipy.ndimage
import os

# Use ffts from scipy 
from scipy import fftpack as FFT
if FFT.fft(np.ones((5,), dtype='single')).dtype.char != 'F':
    print 'Warning: scipy.fftpack is too old: it does not support single precision FFTs.'
   
def gaussian_filter(x, sigma):
    """Complex-valued gaussian_filter
    
    Parameters
    ----------
    x : array_like
        Input array to be filtered
    sigma : float
        Standard deviation of the Gaussian kernel in real space.
        
    See also
    --------
    scipy.ndimage.gaussian_filter : function on which this one is based.
    """
    return scipy.ndimage.gaussian_filter(x.real,sigma) + 1j*scipy.ndimage.gaussian_filter(x.imag,sigma)

FType = None
CType = None

def set_datatype(datatype):
    """Select the floating point precision for reconstructions

    This sets module variables "FType" and "CType".
    
    Parammeters
    -----------
    datatype : {'single', 'double'}
        single/double precision
    """
    global FType, CType
    if datatype not in ['single', 'double']:
        raise RuntimeError("Expected 'single' or 'double'.")
    FType = np.dtype('f' + str(np.dtype(np.typeDict[datatype]).itemsize)).type
    CType = np.dtype('c' + str(2*np.dtype(np.typeDict[datatype]).itemsize)).type
    verbose(3, 'Default data type set to %s' % datatype)

# Start with something
set_datatype('double')

# Shifting functions for subpixel precision in probe positions
def subpix_fourier(a,ramp,forward=True):
    """Perform a sub-pixel shift with the provided phase ramp.
    
    Parameters
    ----------
    a : array_like
        Input array
    ramp : array_like
        Phase ramp to apply
    forward : bool, optional
        Apply the ramp as-is (default). If False, conjugate the ramp first to apply
        the inverse shift.
    """
    if forward:
        return FFT.ifftn(FFT.fftn(a)*ramp)
    else:
        return FFT.ifftn(FFT.fftn(a)*ramp.conj())

def subpix_linear(a, pts, forward=True):
    """Perform a subpixel shift using bilinear interpolation.
    
    Parameters
    ----------
    a : array_like
        Input array (2D)
    pts : 2-tuple of float or array_like
        Shift to apply to the array.
    forward : bool
        Default: apply the given shift. If forward==False, apply the inverse shift.
    """
    x0,x1 = pts
    if not forward: x0,x1 = -x0,-x1
    a00 = np.roll(np.roll(a,-int(np.floor(x0)),axis=0),-int(np.floor(x1)), axis=1)
    a01 = np.roll(np.roll(a,-int(np.floor(x0)),axis=0),-int(np.ceil(x1)), axis=1)
    a10 = np.roll(np.roll(a,-int(np.ceil(x0)),axis=0),-int(np.floor(x1)), axis=1)
    a11 = np.roll(np.roll(a,-int(np.ceil(x0)),axis=0),-int(np.ceil(x1)), axis=1)
    x0 = x0 % 1
    x1 = x1 % 1
    return a00*(1-x0)*(1-x1) + a01*(1-x0)*x1 + a10*x0*(1-x1) + a11*x0*x1

def clean_path(filename):
    """\
    Makes path absolute and create all sub directories if needed. 
    """
    filename = os.path.abspath(os.path.expanduser(filename))
    base = os.path.split(filename)[0]
    if not os.path.exists(base):
        os.makedirs(base)
    return filename

def abs2(x):
    return np.abs(x)**2

    
def make_Fourier_projection(pbound=None, propfct=None):
    """\
    Creates a Fourier projection object.
    Returns the classical version if pbound is None or <=0.
    Otherwise returns the power bound version.
    propfct is a tuple giving the propagation/back-propagation functions,
    default is fftn.
    """
    if pbound is None or pbound <= 0.:
        return Fourier_projection(propfct)
    else:
        return Fourier_projection_pbound(pbound, propfct)

class Fourier_projection(object):
    def __init__(self, propfct=None):
        """\
        Fourier magnitude projection.
        """
        if propfct is None:
            self.fprop = FFT.fftn
            self.ifprop = FFT.ifftn
        else:
            self.fprop, self.ifprop = propfct            
        return
        
    def __call__(self, psi, fmag, fmask=None):
        """\
        Classic Fourier projection - replaces the magnitudes and keeps the phases.
        psi is an iterable over all modes that should sum up to fmag.
        """
        # Fourier transform all modes
        fpsi = [self.fprop(psi) for p in psi]        
        afpsi = np.sqrt(sum(abs2(fp) for fp in fpsi))

        # Compute rescaling of Fourier amplitudes        
        frescale = fmag/(afpsi+1e-10)
        if fmask is not None:
            frescale *= fmask
            frescale += (1-fmask)
            
        psi_proj = [self.ifprop(frescale*fp) for fp in fpsi]
        return psi_proj
        
class Fourier_projection_pbound(object):
    
    def __init__(self, pbound, propfct=None):
        """\
        Fourier magnitude projection with power bound.
        """
        self.pbound = pbound
            
        if propfct is None:
            self.fprop = FFT.fftn
            self.ifprop = FFT.ifftn
        else:
            self.fprop, self.ifprop = propfct
        return
        
    def __call__(self, psi, fmag, fmask=None):
        """\
        Bounded Fourier projection: applies only if the distance to data is larger than self.pbound.
        """
        # Fourier transform all modes
        fpsi = [self.fprop(psi) for p in psi]        
        afpsi = np.sqrt(sum(abs2(fp) for fp in fpsi))

        # Compute deviation
        fdev = afpsi - fmag
        if fmask is None:
            fdev2 = fdev**2
        else:
            fdev2 = fmask*fdev**2
        power = fdev2.mean()

        # If deviation is large enough, rescale amplitudes
        if power > self.pbound:
            renorm = np.sqrt(self.pbound / power)

            # Compute rescaling of Fourier amplitudes        
            if fmask is None:
                frescale = (fmag + fdev*renorm)/(afpsi + 1e-10)
            else:
                frescale = (1-fmask) + fmask*(fmag + fdev*renorm)/(afpsi + 1e-10)

            # Rescale everything and propagate back              
            psi_proj = [self.ifprop(frescale*fp) for fp in fpsi]
        else:
            # Deviation is too small - nothing happens
            psi_proj = psi
            
        return psi_proj
